package codesign

import (
	"archive/tar"
	"compress/gzip"
	"debug/macho"
	"io"
	"io/fs"
	"os"
	"path/filepath"
)

func ResignMacho(path string, asArchive bool, command string, links []string) error {
	// Update Mach-O code signature, if any.
	// After injecting version or image, binaries become invalidated on mac os arm64
	// architectures and we have to sign them according to the new content.
	// Signing is handled exactly the similar way how go handles same situation.
	if !asArchive {
		return signSingleFile(path)
	}

	// archived files are not seekable and go re-sign process needs writerat and readerat
	// interfaces. Therefore, we are extracting these files into temp location, signing there,
	// overriding current path with this new archive file.
	return signTarFile(path, command, links)
}

func signSingleFile(path string) error {
	file, err := os.OpenFile(filepath.Clean(path), os.O_RDWR, 0755)
	if err != nil {
		return err
	}
	defer file.Close()

	wr := io.WriterAt(file)
	if machoFile, cmd, ok := findMachoCodeSignature(wr); ok {
		if Size(int64(cmd.Dataoff), "a.out") == int64(cmd.Datasize) {
			// Update the signature if the size matches, so we don't need to
			// fix up headers. Binaries generated by the Go linker should have
			// the expected size. Otherwise skip.
			text := machoFile.Segment("__TEXT")
			cs := make([]byte, cmd.Datasize)
			Sign(cs, wr.(io.Reader), "a.out", int64(cmd.Dataoff), int64(text.Offset), int64(text.Filesz), machoFile.Type == macho.TypeExec)
			if _, err := wr.WriteAt(cs, int64(cmd.Dataoff)); err != nil {
				return err
			}
		}
	}

	return nil
}

func signTarFile(path, command string, links []string) error {
	tempDir, err := extractTarToTmpAndSign(path)
	if err != nil {
		return err
	}
	defer os.RemoveAll(tempDir)

	outFile, err := os.Create(path)
	if err != nil {
		return err
	}

	gw, err := gzip.NewWriterLevel(outFile, 3)
	if err != nil {
		return err
	}

	tw := tar.NewWriter(gw)
	defer func() {
		tw.Close()
		gw.Close()
	}()

	err = filepath.Walk(tempDir, func(path string, info fs.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		header, err := tar.FileInfoHeader(info, info.Name())
		if err != nil {
			return err
		}

		if err := tw.WriteHeader(header); err != nil {
			return err
		}

		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		_, err = io.Copy(tw, file)
		return err
	})

	if err != nil {
		return err
	}

	for _, l := range links {
		err := tw.WriteHeader(&tar.Header{
			Name:     l,
			Mode:     int64(os.FileMode(0755).Perm()),
			Size:     0,
			Typeflag: tar.TypeLink,
			Linkname: command,
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func extractTarToTmpAndSign(path string) (string, error) {
	tempDir, err := os.MkdirTemp(os.TempDir(), "oc-release-extract-*")
	if err != nil {
		return "", err
	}

	gzipFile, err := os.Open(path)
	if err != nil {
		return "", err
	}

	gr, err := gzip.NewReader(gzipFile)
	if err != nil {
		return "", err
	}
	defer func() {
		gr.Close()
		gzipFile.Close()
	}()

	tw := tar.NewReader(gr)

	for {
		header, err := tw.Next()

		if err == io.EOF {
			break
		}

		if err != nil {
			return "", err
		}

		tempExtractionPath := filepath.Clean(filepath.Join(tempDir, header.Name))

		switch header.Typeflag {
		case tar.TypeReg:
			f, err := os.OpenFile(tempExtractionPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, os.FileMode(header.Mode).Perm())
			if err != nil {
				return "", err
			}

			if _, err := io.Copy(f, tw); err != nil {
				return "", err
			}

			f.Close()

			if err = signSingleFile(tempExtractionPath); err != nil {
				return "", err
			}
		}
	}

	return tempDir, nil
}

func findMachoCodeSignature(r interface{}) (*macho.File, CodeSigCmd, bool) {
	ra, ok := r.(io.ReaderAt)
	if !ok {
		return nil, CodeSigCmd{}, false
	}
	f, err := macho.NewFile(ra)
	if err != nil {
		return nil, CodeSigCmd{}, false
	}
	cmd, ok := FindCodeSigCmd(f)
	return f, cmd, ok
}
